from sqlalchemy import Column, Integer, Index
from sqlalchemy.orm.session import Session

from airflow.models.base import Base
from airflow.utils.session import provide_session


class WorkFlowTask(Base):
    id = Column(Integer, primary_key=True)
    workflow_id = Column(Integer)
    task_id = Column(Integer)

    __tablename__ = "idata_workflow_task"

    __table_args__ = (
        Index('workflow_task_id', workflow_id, task_id),
        Index('task_id', task_id),
    )

    @staticmethod
    @provide_session
    def query(task_id, session: Session = None):
        return session.query(WorkFlowTask).filter(WorkFlowTask.task_id == task_id).all()

    @staticmethod
    @provide_session
    def query_workflow_task(wfid, session: Session = None):
        return session.query(WorkFlowTask).filter(WorkFlowTask.workflow_id == wfid).all()

    @staticmethod
    @provide_session
    def update(workflow_id: int, tasks: list, session: Session = None):
        session.query(WorkFlowTask).filter(WorkFlowTask.workflow_id == workflow_id).delete()
        for tid in tasks:
            dt = WorkFlowTask(workflow_id=workflow_id, task_id=tid)
            session.add(dt)
